import streamlit as st
from typing import List, Dict
from core.api_client import OllamaClient
from core.state_manager import StateManager

class ChatInterface:
    """聊天界面"""
    
    def __init__(self, client: OllamaClient, state_manager: StateManager):
        """初始化聊天界面"""
        self.client = client
        self.state_manager = state_manager
        
        # 初始化消息历史
        if 'messages' not in st.session_state:
            st.session_state.messages = []
        
        # 只在第一次初始化时加载会话
        if 'initialized' not in st.session_state:
            # 获取可用会话
            sessions = self.state_manager.get_sessions()
            if sessions:
                # 加载第一个会话
                first_session = sessions[0]
                if self.state_manager.load_session(first_session):
                    self.client.logger.info(f"已加载默认会话: {first_session}")
                else:
                    self.client.logger.warning(f"加载默认会话失败: {first_session}")
            else:
                # 如果没有会话，创建新会话
                session_id = self.state_manager.new_session()
                self.client.logger.info(f"已创建新会话: {session_id}")
            
            st.session_state.initialized = True
    
    def _format_message(self, content: str) -> str:
        """格式化消息内容，处理代码块"""
        # 查找所有代码块
        import re
        code_blocks = re.finditer(r'```(\w+)?\n(.*?)```', content, re.DOTALL)
        formatted_parts = []
        last_end = 0
        
        for match in code_blocks:
            # 添加代码块前的文本
            if match.start() > last_end:
                formatted_parts.append(content[last_end:match.start()])
            
            # 获取语言和代码
            lang = match.group(1) or ''
            code = match.group(2).strip()
            
            # 格式化代码块
            formatted_parts.append(f'<pre><code class="language-{lang}">{code}</code></pre>')
            last_end = match.end()
        
        # 添加剩余文本
        if last_end < len(content):
            formatted_parts.append(content[last_end:])
        
        return ''.join(formatted_parts)
    
    def _display_message(self, message: Dict):
        """显示单条消息"""
        with st.chat_message(message["role"]):
            content = message["content"]
            
            # 查找所有代码块
            import re
            code_blocks = re.finditer(r'```(\w+)?\n(.*?)```', content, re.DOTALL)
            last_end = 0
            
            for match in code_blocks:
                # 显示代码块前的文本
                if match.start() > last_end:
                    st.markdown(content[last_end:match.start()])
                
                # 获取语言和代码
                lang = match.group(1) or ''
                code = match.group(2).strip()
                
                # 创建代码块容器
                col1, col2 = st.columns([10, 1])
                with col1:
                    st.code(code, language=lang, line_numbers=True)
                with col2:
                    if st.button("📋", key=f"copy_{hash(code)}", help="复制代码"):
                        st.markdown(f"""
                            <script>
                                navigator.clipboard.writeText(`{code}`);
                            </script>
                        """, unsafe_allow_html=True)
                        st.toast("代码已复制！")
                
                last_end = match.end()
            
            # 显示剩余文本
            if last_end < len(content):
                st.markdown(content[last_end:])
    
    def _display_messages(self):
        """显示所有消息"""
        for message in st.session_state.messages:
            self._display_message(message)
    
    def _handle_user_input(self):
        """处理用户输入"""
        if prompt := st.chat_input("输入消息..."):
            # 添加用户消息
            user_message = {
                "role": "user",
                "content": prompt
            }
            st.session_state.messages.append(user_message)
            self._display_message(user_message)
            
            # 创建AI响应占位符
            with st.chat_message("assistant"):
                message_placeholder = st.empty()
                full_response = ""
                
                try:
                    # 准备完整的消息历史，包括系统提示
                    messages = [
                        {
                            "role": "system",
                            "content": "你是一个有帮助的AI助手。请用简洁、专业的方式回答问题。"
                        }
                    ]
                    
                    # 添加历史对话
                    messages.extend([
                        {"role": m["role"], "content": m["content"]} 
                        for m in st.session_state.messages[:-1]  # 不包括最新的用户消息
                    ])
                    
                    # 添加最新的用户消息
                    messages.append({
                        "role": "user",
                        "content": prompt
                    })
                    
                    # 获取流式响应
                    for chunk in self.client.chat_stream(
                        messages=messages,
                        model=self.state_manager.get_state("model", "llama2"),
                        temperature=self.state_manager.get_state("temperature", 0.7),
                        max_length=self.state_manager.get_state("max_length", 2000)
                    ):
                        if chunk and "response" in chunk:
                            full_response += chunk["response"]
                            # 实时更新显示，使用格式化后的内容
                            message_placeholder.markdown(self._format_message(full_response) + "▌")
                    
                    # 显示完整响应
                    message_placeholder.markdown(self._format_message(full_response))
                    
                    # 保存响应
                    assistant_message = {
                        "role": "assistant",
                        "content": full_response
                    }
                    st.session_state.messages.append(assistant_message)
                    
                    # 保存会话
                    if self.state_manager.get_state("current_session"):
                        self.state_manager._save_session()
                    
                except ValueError as e:
                    st.error(f"API格式错误: {str(e)}")
                    self.client.logger.error(f"API格式错误: {str(e)}")
                except ConnectionError as e:
                    st.error(f"连接错误: {str(e)}")
                    self.client.logger.error(f"连接错误: {str(e)}")
                except Exception as e:
                    st.error(f"发生错误: {str(e)}")
                    self.client.logger.error(f"聊天错误: {str(e)}")
    
    def _show_sidebar(self):
        """显示侧边栏"""
        with st.sidebar:
            st.title("设置")
            
            # 思索模式开关
            thinking_mode = st.toggle(
                "显示思考过程",
                value=False,
                help="开启后会显示AI的思考过程"
            )
            self.client.set_thinking_mode(thinking_mode)
            
            # 获取可用模型列表
            try:
                models = self.client.list_models()
                model_names = [m["name"] for m in models] or ["llama2"]
                
                # 模型选择
                model = st.selectbox(
                    "选择模型",
                    options=model_names,
                    index=0,
                    help="选择要使用的模型"
                )
                
                # 显示模型信息
                if models:
                    selected_model = next((m for m in models if m["name"] == model), None)
                    if selected_model:
                        st.caption(f"大小: {selected_model['size']}")
                        st.caption(f"更新时间: {selected_model['modified_at']}")
            except Exception as e:
                st.error(f"获取模型列表失败: {str(e)}")
                model = "llama2"
            
            # 温度设置
            temperature = st.slider(
                "温度",
                min_value=0.0,
                max_value=1.0,
                value=0.7,
                step=0.1,
                help="控制响应的随机性，值越高随机性越大"
            )
            
            # 最大长度设置
            max_length = st.number_input(
                "最大长度",
                min_value=100,
                max_value=4096,
                value=2000,
                help="生成文本的最大长度"
            )
            
            # 更新设置
            self.state_manager.update_states({
                "model": model,
                "temperature": temperature,
                "max_length": max_length
            })
            
            # 会话管理
            st.divider()
            st.subheader("会话管理")
            
            # 会话列表
            sessions = self.state_manager.get_sessions()
            if sessions:
                current_session = self.state_manager.get_current_session()
                
                # 选择会话
                selected_session = st.selectbox(
                    "选择会话",
                    options=sessions,
                    index=sessions.index(current_session) if current_session in sessions else 0,
                    format_func=lambda x: f"会话 {x}",
                    key="session_selector"
                )
                
                # 自动切换会话
                if selected_session != current_session:
                    if self.state_manager.load_session(selected_session):
                        st.session_state.messages = self.state_manager.get_messages().copy()
                        st.rerun()
                
                # 删除按钮
                if st.button("删除会话", key="delete_session"):
                    if self.state_manager.delete_session(selected_session):
                        if selected_session == current_session:
                            st.session_state.messages = []
                        st.success("删除成功")
                        st.rerun()
                    else:
                        st.error("删除失败")
            
            # 新建会话按钮
            if st.button("新建会话", key="new_session"):
                session_id = self.state_manager.new_session()
                st.session_state.messages = []
                st.success(f"已创建新会话")
                st.rerun()
            
            # 清空当前会话按钮
            if st.button("清空当前会话", key="clear_session"):
                self.state_manager.clear_messages()
                st.session_state.messages = []
                st.rerun()
    
    def render(self):
        """渲染聊天界面"""
        self._show_sidebar()
        
        st.title("Ollama Chat")
        
        # 显示消息历史
        self._display_messages()
        
        # 处理用户输入
        self._handle_user_input() 